#include "gpu/spatical_sampling.h"
#include "gpu/kernel.h"

inline void checkCudaError(cudaError_t result, const char* message) {
    if (result != cudaSuccess) {
        std::cerr << "CUDA error: " << message << " - " << cudaGetErrorString(result) << std::endl;
        exit(EXIT_FAILURE);
    }
}

SpaticalSampler::SpaticalSampler() {
    checkCudaError(cudaStreamCreate(&stream_), "cudaStreamCreate");
    int memory_size = 1024 * 1024 * 300; // 300M 内存池
    buffer_ = static_cast<char*>(aligned_alloc(kAligned, memory_size)); 
    cpu_buffer_ = static_cast<char*>(aligned_alloc(kAligned, memory_size)); 
    memset(buffer_, 0, memory_size);
    memset(cpu_buffer_, 0, memory_size);

    cudaMalloc(&gpu_buffer_, memory_size);
    cudaMemset(gpu_buffer_, 0, memory_size);
}

SpaticalSampler::~SpaticalSampler() { 
    free(buffer_); 
    free(cpu_buffer_);
    cudaFree(gpu_buffer_);
}

int SpaticalSampler::init_memory(const std::vector<RefPoint>& refline, const std::vector<Path>& path_vec) {
    int refine_data_memory_size = refline.size() * sizeof(RefPoint);
    int path_nums = path_vec.size();

    // 1. DataInfo 序列化到memroy中
    DataInfo* data_info = reinterpret_cast<DataInfo*>(buffer_);
    data_info->path_nums = path_nums;
    data_info->min_ref_s = refline.front().s;
    data_info->max_ref_s = refline.back().s;
    data_info->refline_offset = MemoryAligned(buffer_ + sizeof(DataInfo), kAligned) - reinterpret_cast<uintptr_t>(buffer_);
    data_info->path_info_offset = MemoryAligned(buffer_ + sizeof(DataInfo) + refine_data_memory_size, kAligned) - reinterpret_cast<uintptr_t>(buffer_);
    data_info->path_coeffs_offset =
        MemoryAligned(buffer_ + sizeof(DataInfo) + refine_data_memory_size + sizeof(PathInfo) * path_nums, kAligned) - reinterpret_cast<uintptr_t>(buffer_);

    // 2. RefPoint 序列化到memroy中
    RefPoint* ref_point_start_ptr = reinterpret_cast<RefPoint*>(buffer_ + data_info->refline_offset);
    memcpy((void*)ref_point_start_ptr, (void*)refline.data(), refine_data_memory_size);

    // 2. PathInfo 序列化到memroy中
    PathInfo* path_info_start_ptr = reinterpret_cast<PathInfo*>(buffer_ + data_info->path_info_offset);
    PathInfo* path_info = path_info_start_ptr;
    for (int i = 0; i < path_vec.size(); i++) {
        int size = path_vec[i].serialize_path_info(path_info);
        path_info++;
    }
    int path_info_memory_size = reinterpret_cast<char*>(path_info) - reinterpret_cast<char*>(path_info_start_ptr);
    std::cout << "path info, used_memory_size: " << path_info_memory_size / 1024.0 << "kB" << std::endl;

    // 3. CubicPolynomialCoeffs 序列化到memroy中
    CubicPolynomialCoeffs* path_coeffs_start_ptr = reinterpret_cast<CubicPolynomialCoeffs*>(buffer_ + data_info->path_coeffs_offset);
    CubicPolynomialCoeffs* path_coeffs = path_coeffs_start_ptr;
    for (int i = 0; i < path_vec.size(); i++) {
        int coeff_nums = path_vec[i].serialize_path_coeffs(path_coeffs);
        PathInfo* curr_path_info = path_info_start_ptr + i;
        curr_path_info->cubic_polynomial_coeff_nums = coeff_nums;
        curr_path_info->cubic_polynomial_coeff_offset = reinterpret_cast<char*>(path_coeffs) - reinterpret_cast<char*>(buffer_);
        // std::cout << "path_coeffs ptr: " << path_coeffs << ", cubic_polynomial_coeff_offset: " << curr_path_info->cubic_polynomial_coeff_offset <<
        // std::endl;
        path_coeffs += coeff_nums;
    }
    int path_coeffs_memory_size = reinterpret_cast<char*>(path_coeffs) - reinterpret_cast<char*>(path_coeffs_start_ptr);
    std::cout << "path coeffs, used_memory_size: " << path_coeffs_memory_size / 1024.0 << "kB" << std::endl;

    // 4. 将PointInfo的一些描述信息写到 PathInfo所在内存中
    data_info->point_info_offset = MemoryAligned(path_coeffs, kAligned) - reinterpret_cast<uintptr_t>(buffer_);
    PointInfo* point_info_start_ptr = reinterpret_cast<PointInfo*>(buffer_ + data_info->point_info_offset);
    PointInfo* point_info = point_info_start_ptr;
    for (int i = 0; i < path_vec.size(); i++) {
        PathInfo* curr_path_info = path_info_start_ptr + i;
        int need_point_info_nums = static_cast<int>((curr_path_info->max_s -  curr_path_info->min_s + 1e-4) / kPointInfoResolution);
        curr_path_info->point_info_nums = need_point_info_nums;
        curr_path_info->point_info_offset = reinterpret_cast<char*>(point_info) - reinterpret_cast<char*>(buffer_);
        point_info += need_point_info_nums;
    }

    float max_length = 0;
    for(int i=0; i< path_vec.size(); i++){
        max_length = std::max(path_vec[i].length(), max_length);
    }
    data_info->max_length = max_length;

    std::cout << "data_info, path_nums: " << data_info->path_nums 
              << ", path_info_offset: " << data_info->path_info_offset
              << ", path_coeffs_offset: " << data_info->path_coeffs_offset 
              << ", point_info_offset: " << data_info->point_info_offset 
              << std::endl;

    return reinterpret_cast<char*>(point_info) - reinterpret_cast<char*>(buffer_);
}

int SpaticalSampler::sampling_on_gpu(int size) {
    // 异步过程，隐藏拷贝耗时
    DataInfo* data_info = reinterpret_cast<DataInfo*>(buffer_);
    int info_size = data_info->point_info_offset;
    checkCudaError(cudaMemcpyAsync(gpu_buffer_, buffer_, info_size, cudaMemcpyHostToDevice, stream_), "cudaMemcpyAsync (HostToDevice)");

    // 开始计算
    spatical_sampling_launcher(&stream_, data_info, gpu_buffer_);

    // 异步过程，隐藏拷贝耗时
    checkCudaError(cudaMemcpyAsync(cpu_buffer_, gpu_buffer_, size, cudaMemcpyDeviceToHost, stream_), "cudaMemcpyAsync (DeviceToHost)");  
    
    // 同步
    checkCudaError(cudaStreamSynchronize(stream_), "cudaStreamSynchronize");
    return 0;
}

void SpaticalSampler::debug_print(bool is_gpu) {
    char* data = is_gpu ? cpu_buffer_ : buffer_;

    DataInfo* data_info = reinterpret_cast<DataInfo*>(data);
    int path_nums = data_info->path_nums;
    std::cout << "data_info, path_nums: " << data_info->path_nums 
              << ", path_info_offset: " << data_info->path_info_offset
              << ", path_coeffs_offset: " << data_info->path_coeffs_offset 
              << ", point_info_offset: " << data_info->point_info_offset 
              << std::endl;

    PathInfo* path_info_start_ptr = reinterpret_cast<PathInfo*>(data + data_info->path_info_offset);
    PathInfo* path_info = path_info_start_ptr;
    for (int i = 0; i < path_nums; i++) {
        int cubic_polynomial_coeff_nums = path_info->cubic_polynomial_coeff_nums;
        CubicPolynomialCoeffs* path_coeffs = reinterpret_cast<CubicPolynomialCoeffs*>(data + path_info->cubic_polynomial_coeff_offset);

        std::cout << "ptr: " << path_info 
                  << ", path index: " << path_info->path_index 
                  << ", min_s: " << path_info->min_s 
                  << ", max_s: " << path_info->max_s
                  << ", cubic_polynomial_coeff_nums: " << path_info->cubic_polynomial_coeff_nums
                  << ", cubic_polynomial_coeff_offset: " << path_info->cubic_polynomial_coeff_offset 
                  << ", point_info_nums: " << path_info->point_info_nums
                  << ", point_info_offset: " << path_info->point_info_offset
                  << std::endl;

        for (int j = 0; j < cubic_polynomial_coeff_nums; j++) {
            std::cout << "ptr: " << path_coeffs << ", coeffs.index: " << j << ", coeffs.start_s: " << path_coeffs->start_s
                      << ", coeffs.end_s: " << path_coeffs->end_s << ", coeffs.a3: " << path_coeffs->a3 << std::endl;
            path_coeffs++;
        }

        path_info++;
    }
}

void SpaticalSampler::dump(std::ofstream& out_file, bool is_gpu) const {
    char* data = is_gpu ? cpu_buffer_ : buffer_;
    DataInfo* data_info = reinterpret_cast<DataInfo*>(data);
    int path_nums = data_info->path_nums;

    PathInfo* path_info_start_ptr = reinterpret_cast<PathInfo*>(data + data_info->path_info_offset);
    PathInfo* path_info = path_info_start_ptr;

    for (int i = 0; i < path_nums; i++) {
        int path_index = path_info->path_index;
        int point_info_nums = path_info->point_info_nums;
        int point_info_offset = path_info->point_info_offset;
        
        out_file << "path_index: "<< path_index << ", point_info_offset: " << point_info_offset << ", ";
        PointInfo* point_info = reinterpret_cast<PointInfo*>(data + point_info_offset);
        for(int i=0; i< point_info_nums; i++){
            if(std::abs(point_info->s) < 200  && std::abs(point_info->l) < 20){ // trick
                out_file << "(";
                out_file << point_info->x << ", " << point_info->y;
                out_file << "), ";
                point_info++;
            }else{
                std::cout << "error" << ", path_index: "<< path_index 
                          << ", point_info_offset: " << point_info_offset 
                          << ", point_index: "<< i 
                          << ", s: "<<  point_info->s 
                          << ", l: "<<  point_info->l 
                          << std::endl;
            }
        }
        out_file << std::endl;
        path_info++;
    }
}


